static BASE64_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

pub fn encode(bytes: &[u8]) -> String {
    bytes
        .chunks(3)
        .flat_map(|chunk| {
            let mut n = chunk.iter().fold(0usize, |n, &b| (n << 8) | b as usize);
            let padding = 3 - chunk.len();
            n <<= padding * 8;
            (0..4).rev().map(move |j| {
                if j >= padding {
                    BASE64_CHARS[(n >> j * 6) & 0x3F] as char
                } else {
                    '='
                }
            })
        })
        .collect()
}

pub fn decode(base64: &str) -> Vec<u8> {
    let mut result = Vec::new();
    let mut n = 0;
    let mut bits = 0;
    for c in base64.bytes() {
        let index = match c {
            b'A'..=b'Z' => c - b'A',
            b'a'..=b'z' => c - b'a' + 26,
            b'0'..=b'9' => c - b'0' + 26 + 26,
            b'+' => 62,
            b'/' => 63,
            b'=' => continue,
            _ => panic!("Invalid base64 character"),
        };
        n = (n << 6) | index as u32;

        bits += 6;
        if bits >= 8 {
            bits -= 8;
            result.push((n >> bits) as u8);
            n &= (1 << bits) - 1;
        }
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[track_caller]
    fn test_case(input: &[u8], expected: &str) {
        assert_eq!(encode(input), expected);
        assert_eq!(decode(expected), input);
    }

    #[track_caller]
    fn test_case_rev(input: &str, expected: &[u8]) {
        assert_eq!(decode(input), expected);
        assert_eq!(encode(expected), input);
    }

    #[test]
    fn test_encode_decode() {
        test_case(b"\xfd\x06\x5b\xed\x2d\x09\x5a\x9f", "/QZb7S0JWp8=");
    }

    #[test]
    fn test_decode_encode() {
        test_case_rev("/QZb7S0JWp8=", b"\xfd\x06\x5b\xed\x2d\x09\x5a\x9f");
    }
}
